#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Hugging Face镜像源配置脚本
解决国内网络环境下模型下载问题
"""

import os
import sys
import subprocess
import platform

def setup_hf_mirror():
    """配置Hugging Face镜像源"""
    print("🔧 配置Hugging Face镜像源...")
    
    # 设置环境变量
    env_vars = {
        'HF_ENDPOINT': 'https://hf-mirror.com',
        'HF_HUB_URL': 'https://hf-mirror.com',
        'HF_HOME': os.path.expanduser('~/.cache/huggingface'),
        'TRANSFORMERS_CACHE': os.path.expanduser('~/.cache/huggingface/transformers'),
        'HF_DATASETS_CACHE': os.path.expanduser('~/.cache/huggingface/datasets'),
        'HF_HUB_DISABLE_TELEMETRY': '1',  # 禁用遥测
        'HF_HUB_DISABLE_IMPLICIT_TOKEN': '1',  # 禁用隐式token
    }
    
    # 设置环境变量
    for key, value in env_vars.items():
        os.environ[key] = value
        print(f"   {key} = {value}")
    
    # 创建缓存目录
    cache_dir = env_vars['HF_HOME']
    os.makedirs(cache_dir, exist_ok=True)
    print(f"✅ 缓存目录创建成功: {cache_dir}")
    
    return env_vars

def test_hf_connection():
    """测试Hugging Face连接"""
    print("\n🔍 测试Hugging Face连接...")
    
    try:
        import requests
        
        # 测试镜像源连接
        test_url = "https://hf-mirror.com/api/models"
        response = requests.get(test_url, timeout=10)
        
        if response.status_code == 200:
            print("✅ Hugging Face镜像源连接正常")
            return True
        else:
            print(f"❌ 镜像源连接异常，状态码: {response.status_code}")
            return False
            
    except Exception as e:
        print(f"❌ 连接测试失败: {e}")
        return False

def download_model_locally():
    """本地下载模型（可选）"""
    print("\n📥 本地模型下载选项:")
    print("1. 使用镜像源自动下载（推荐）")
    print("2. 手动下载模型文件")
    print("3. 跳过下载，使用备选方案")
    
    choice = input("请选择 (1-3): ").strip()
    
    if choice == "1":
        print("✅ 将使用镜像源自动下载模型")
        return "auto"
    elif choice == "2":
        print("📋 手动下载说明:")
        print("1. 访问 https://hf-mirror.com")
        print("2. 搜索并下载以下模型:")
        print("   - shibing624/text2vec-base-chinese")
        print("   - GanymedeNil/text2vec-large-chinese")
        print("3. 将模型文件放置到缓存目录")
        return "manual"
    else:
        print("✅ 将使用备选方案")
        return "fallback"

def create_model_config():
    """创建模型配置文件"""
    config_content = """# Hugging Face模型配置
# 此文件用于配置模型下载和缓存

# 镜像源配置
HF_ENDPOINT=https://hf-mirror.com
HF_HUB_URL=https://hf-mirror.com

# 缓存目录配置
HF_HOME=~/.cache/huggingface
TRANSFORMERS_CACHE=~/.cache/huggingface/transformers
HF_DATASETS_CACHE=~/.cache/huggingface/datasets

# 其他配置
HF_HUB_DISABLE_TELEMETRY=1
HF_HUB_DISABLE_IMPLICIT_TOKEN=1

# 模型优先级列表
MODEL_PRIORITY=[
    "shibing624/text2vec-base-chinese",
    "GanymedeNil/text2vec-large-chinese", 
    "paraphrase-multilingual-MiniLM-L12-v2",
    "all-MiniLM-L6-v2"
]
"""
    
    config_file = "hf_config.env"
    with open(config_file, 'w', encoding='utf-8') as f:
        f.write(config_content)
    
    print(f"✅ 配置文件创建成功: {config_file}")
    return config_file

def install_requirements():
    """安装必要的依赖"""
    print("\n📦 检查依赖包...")
    
    required_packages = [
        'sentence-transformers',
        'torch',
        'transformers',
        'huggingface_hub'
    ]
    
    missing_packages = []
    
    for package in required_packages:
        try:
            __import__(package.replace('-', '_'))
            print(f"✅ {package} 已安装")
        except ImportError:
            print(f"❌ {package} 未安装")
            missing_packages.append(package)
    
    if missing_packages:
        print(f"\n📥 需要安装的包: {', '.join(missing_packages)}")
        install = input("是否自动安装? (y/n): ").strip().lower()
        
        if install == 'y':
            try:
                subprocess.check_call([
                    sys.executable, '-m', 'pip', 'install'
                ] + missing_packages)
                print("✅ 依赖包安装完成")
            except subprocess.CalledProcessError as e:
                print(f"❌ 安装失败: {e}")
                print("请手动安装: pip install " + " ".join(missing_packages))
        else:
            print("请手动安装依赖包")
    
    return len(missing_packages) == 0

def create_startup_script():
    """创建启动脚本"""
    if platform.system() == "Windows":
        script_content = """@echo off
REM SecurityWarning项目启动脚本
echo 正在启动SecurityWarning项目...

REM 设置Hugging Face镜像源
set HF_ENDPOINT=https://hf-mirror.com
set HF_HUB_URL=https://hf-mirror.com
set HF_HOME=%USERPROFILE%\\.cache\\huggingface
set TRANSFORMERS_CACHE=%USERPROFILE%\\.cache\\huggingface\\transformers
set HF_DATASETS_CACHE=%USERPROFILE%\\.cache\\huggingface\\datasets
set HF_HUB_DISABLE_TELEMETRY=1
set HF_HUB_DISABLE_IMPLICIT_TOKEN=1

REM 启动项目
python run.py
pause
"""
        script_file = "start.bat"
    else:
        script_content = """#!/bin/bash
# SecurityWarning项目启动脚本
echo "正在启动SecurityWarning项目..."

# 设置Hugging Face镜像源
export HF_ENDPOINT="https://hf-mirror.com"
export HF_HUB_URL="https://hf-mirror.com"
export HF_HOME="$HOME/.cache/huggingface"
export TRANSFORMERS_CACHE="$HOME/.cache/huggingface/transformers"
export HF_DATASETS_CACHE="$HOME/.cache/huggingface/datasets"
export HF_HUB_DISABLE_TELEMETRY="1"
export HF_HUB_DISABLE_IMPLICIT_TOKEN="1"

# 启动项目
python run.py
"""
        script_file = "start.sh"
        # 设置执行权限
        os.chmod(script_file, 0o755)
    
    with open(script_file, 'w', encoding='utf-8') as f:
        f.write(script_content)
    
    print(f"✅ 启动脚本创建成功: {script_file}")
    return script_file

def main():
    """主函数"""
    print("🚀 Hugging Face镜像源配置工具")
    print("=" * 50)
    
    # 1. 设置环境变量
    env_vars = setup_hf_mirror()
    
    # 2. 测试连接
    connection_ok = test_hf_connection()
    
    # 3. 检查依赖
    deps_ok = install_requirements()
    
    # 4. 创建配置文件
    config_file = create_model_config()
    
    # 5. 创建启动脚本
    startup_script = create_startup_script()
    
    # 6. 下载选项
    download_mode = download_model_locally()
    
    print("\n" + "=" * 50)
    print("📋 配置总结:")
    print(f"   ✅ 环境变量已设置")
    print(f"   {'✅' if connection_ok else '❌'} 连接测试: {'成功' if connection_ok else '失败'}")
    print(f"   {'✅' if deps_ok else '❌'} 依赖检查: {'通过' if deps_ok else '需要安装'}")
    print(f"   ✅ 配置文件: {config_file}")
    print(f"   ✅ 启动脚本: {startup_script}")
    print(f"   📥 下载模式: {download_mode}")
    
    print("\n🚀 使用方法:")
    if platform.system() == "Windows":
        print(f"   1. 双击运行: {startup_script}")
        print(f"   2. 或在命令行运行: python run.py")
    else:
        print(f"   1. 运行启动脚本: ./{startup_script}")
        print(f"   2. 或在命令行运行: python run.py")
    
    print("\n💡 注意事项:")
    print("   - 首次运行可能需要下载模型，请耐心等待")
    print("   - 如果下载失败，系统会自动使用备选方案")
    print("   - 模型文件会缓存在 ~/.cache/huggingface 目录")
    
    print("\n✅ 配置完成！")

if __name__ == "__main__":
    main() 